#!/usr/bin/env python3
# Author: Armit
# Create Time: 2024/04/07

# view cossim of trainable embeddings

from pathlib import Path
from argparse import ArgumentParser

import torch
import torch.nn.functional as F
from torch import Tensor
import matplotlib.pyplot as plt


def cos_sim(a: Tensor, b: Tensor) -> Tensor:
  if not isinstance(a, Tensor): a = torch.tensor(a)
  if not isinstance(b, Tensor): b = torch.tensor(b)
  if len(a.shape) == 1: a = a.unsqueeze(0)
  if len(b.shape) == 1: b = b.unsqueeze(0)
  a_norm = F.normalize(a, p=2, dim=1)
  b_norm = F.normalize(b, p=2, dim=1)
  return torch.mm(a_norm, b_norm.transpose(0, 1))


def dot_score(a: Tensor, b: Tensor) -> Tensor:
  if not isinstance(a, Tensor): a = torch.tensor(a)
  if not isinstance(b, Tensor): b = torch.tensor(b)
  if len(a.shape) == 1: a = a.unsqueeze(0)
  if len(b.shape) == 1: b = b.unsqueeze(0)
  return torch.mm(a, b.transpose(0, 1))


def plot2(cosim:Tensor, dotsc:Tensor, title:str=''):
  plt.clf()
  plt.subplot(121) ; plt.imshow(cosim.cpu().numpy()) ; plt.title('cosim') ; plt.gca().invert_yaxis()
  plt.subplot(122) ; plt.imshow(dotsc.cpu().numpy()) ; plt.title('dotsc') ; plt.gca().invert_yaxis()
  plt.suptitle(title)
  plt.show()
  plt.close()


@torch.inference_mode()
def run(args):
  ckpt = torch.load(args.f, map_location='cpu')

  if 'appearance' in ckpt:
    E: Tensor = ckpt['appearance']['embedding.weight']
    print('appearance_embed.shape:', E.shape)
    cosim = cos_sim(E, E)
    dotsc = dot_score(E, E)
    plot2(cosim, dotsc, 'appearance')

  if 'occlusion' in ckpt:
    E: Tensor = ckpt['occlusion']['embedding.weight']
    print('occlusion_embed.shape:', E.shape)
    cosim = cos_sim(E, E)
    dotsc = dot_score(E, E)
    plot2(cosim, dotsc, 'occlusion')


if __name__ == '__main__':
  paser = ArgumentParser()
  paser.add_argument('-f', type=Path, help='path to *.ckpt file')
  args = paser.parse_args()

  run(args)
